#include "../Base.h"

VOID
DefaultCreateProceeNotify(
    _In_ HANDLE ParentId,
    _In_ HANDLE ProcessId,
    _In_ BOOLEAN Create
) {

    PEPROCESS process;
    NTSTATUS status = PsLookupProcessByProcessId(ProcessId, &process);
    if (!NT_SUCCESS(status)) {
        LOG("查找进程失败%s", __FUNCTION__);
        return;
    }
    //创建进程
    if (Create) {
     

        LOG("创建进程: %s", PsGetProcessImageFileName(process));
        
    }
    //进程死亡
    else {
        LOG("销毁进程: %s", PsGetProcessImageFileName(process));
    }
    ObDereferenceObject(process);
}

NTSTATUS SetCreateProcessCallback(PCREATE_PROCESS_NOTIFY_ROUTINE notify) {
    NTSTATUS status;
    if (!notify) {
        status = PsSetCreateProcessNotifyRoutine(DefaultCreateProceeNotify, FALSE);
    }
    else {
        status = PsSetCreateProcessNotifyRoutine(notify, FALSE);
    }
 
    if (!NT_SUCCESS(status)) {
        LOG("进程回调创建失败!");
    }
    return status;
}


VOID
DefaultCreateThreadCallback(
    _In_ HANDLE ProcessId,
    _In_ HANDLE ThreadId,
    _In_ BOOLEAN Create
) {
    if (Create) {
        LOG("线程启动！id = %d", ThreadId);
    }
    else{
        LOG("线程销毁！id = %d", ThreadId);
    }

}

NTSTATUS SetCreateThreadCallback(PCREATE_THREAD_NOTIFY_ROUTINE notify) {
    NTSTATUS  status;
    if (!notify) {
        status = PsSetCreateThreadNotifyRoutine(DefaultCreateThreadCallback);
    }
    else {
        status = PsSetCreateThreadNotifyRoutine(notify);
       
    }

    if (!NT_SUCCESS(status)) {
        LOG("线程回调创建失败!");
    }
    return status;
}


VOID
DefaultLoadImageCallback(
    _In_opt_ PUNICODE_STRING FullImageName,
    _In_ HANDLE ProcessId,                // pid into which image is being mapped
    _In_ PIMAGE_INFO ImageInfo
){

    LOG("加载%ws", ( *FullImageName).Buffer);
}

NTSTATUS SetLoadImageCallback(PLOAD_IMAGE_NOTIFY_ROUTINE notify) {
    NTSTATUS  status;
    if (!notify) {
        status = PsSetLoadImageNotifyRoutine(DefaultLoadImageCallback);
    }
    else {
        status = PsSetLoadImageNotifyRoutine(notify);
    }

    if (!NT_SUCCESS(status)) {
        LOG("模块加载回调创建失败!");
    }
    return status;
}



PULONG GetPspNotifyEnableMask() {




    UNICODE_STRING unName = { 0 };
    RtlInitUnicodeString(&unName, L"PsSetLoadImageNotifyRoutineEx");
    PUCHAR func = (PUCHAR)MmGetSystemRoutineAddress(&unName);
    if (!func) {
        RtlInitUnicodeString(&unName, L"PsSetLoadImageNotifyRoutine");
        func = (PUCHAR)MmGetSystemRoutineAddress(&unName);
        if (!func) return 0;
   }

   

    LONG offset = 0;
    PULONG address = NULL;
    for (int i = 0; i < 0x100; i++) {

        if (func[i] == 0xf0
            && func[i + 1] == 0xf
            && func[i + 2] == 0xBA
            && func[i + 3] == 0x2D) {
            offset = *(PULONG)(func + i + 4);
            ULONG64 next = (ULONG64)(func + i + 9);
            address = (PULONG)(next + offset);
        }
    }



    return address;


}

ULONG PatchNotificationMask(PULONG NotifyEnableMask) {

    if (MmIsAddressValid(NotifyEnableMask)) {
        ULONG oldValue = *NotifyEnableMask;
        *NotifyEnableMask = 0;
        return oldValue;
    }

    return 0;

}


VOID RePatchNotificationMask(PULONG NotifyEnableMask,ULONG oldValue) {

    if (MmIsAddressValid(NotifyEnableMask)) {
       
        *NotifyEnableMask = oldValue;
       
    }
}